package yuekao9.util;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import org.apache.flink.api.common.state.BroadcastState;
import org.apache.flink.api.common.state.MapStateDescriptor;
import org.apache.flink.api.common.state.ReadOnlyBroadcastState;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.co.BroadcastProcessFunction;
import org.apache.flink.util.Collector;

import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.PreparedStatement;

public class BroadcastConnect extends BroadcastProcessFunction<String, String, String> {
    MapStateDescriptor<String, String> mapStateDescriptor;
    Connection conn = null;

    public BroadcastConnect(MapStateDescriptor<String, String> mapStateDescriptor) {
        mapStateDescriptor = this.mapStateDescriptor;
    }

    @Override
    public void open(Configuration parameters) throws Exception {
        Class.forName("org.apache.phoenix.jdbc.PhoenixDriver");
        conn = DriverManager.getConnection("jdbc:phoenix:hadoop-single:2181", "", "");
        conn.setAutoCommit(true);
    }

    @Override
    public void close() throws Exception {
        conn.close();
    }

    @Override
    public void processElement(String s, BroadcastProcessFunction<String, String, String>.ReadOnlyContext readOnlyContext, Collector<String> collector) throws Exception {
        ReadOnlyBroadcastState<String, String> broadcastState = readOnlyContext.getBroadcastState(mapStateDescriptor);

        JSONObject jsonObject = JSON.parseObject(s);
        String table = jsonObject.getString("table");
        String tables = broadcastState.get(table);

        if (tables != null) {
            PreparedStatement ps = conn.prepareStatement("upsert into " + tables + " values (?.?)");
            ps.setObject(1, jsonObject.getJSONObject("data").getInteger("id"));
            ps.setObject(2, jsonObject.getJSONObject("data").toJSONString());
            ps.executeUpdate();
            ps.close();
        }
    }

    @Override
    public void processBroadcastElement(String s, BroadcastProcessFunction<String, String, String>.Context context, Collector<String> collector) throws Exception {
        JSONObject jsonObject = JSON.parseObject(s).getJSONObject("data");
        String source_table = jsonObject.getString("source_table");
        String sink_table = jsonObject.getString("sink_table");

        BroadcastState<String, String> broadcastState = context.getBroadcastState(mapStateDescriptor);
        broadcastState.put(source_table, sink_table);
//create table dim_hospital(id integer PRIMARY KEY, line varchar);
        PreparedStatement ps = conn.prepareStatement("create table if not exists " + sink_table + " (id integer PRIMARY KEY, line varchar)");
        ps.executeUpdate();
    }
}
